# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
import argparse
import datetime
import numpy as np
import time
import torch
import torch.backends.cudnn as cudnn
import json

from pathlib import Path

from timm.data import Mixup
from timm.models import create_model
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.scheduler import create_scheduler
from timm.optim import create_optimizer
from timm.utils import NativeScaler, get_state_dict, ModelEma


from datasets import build_dataset
from engine import train_one_epoch, evaluate
from losses import DistillationLoss
from samplers import RASampler
import models
import utils


def get_args_parser():
    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
    parser.add_argument('--batch-size', default=64, type=int)
    parser.add_argument('--epochs', default=300, type=int)

    # Model parameters
    parser.add_argument('--model', default='deit_small_patch16_224', type=str, metavar='MODEL',
                        help='Name of model to train')
    parser.add_argument('--input-size', default=224, type=int, help='images input size')

    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
                        help='Dropout rate (default: 0.)')
    parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT',
                        help='Drop path rate (default: 0.1)')

    parser.add_argument('--model-ema', action='store_true')
    parser.add_argument('--no-model-ema', action='store_false', dest='model_ema')
    parser.set_defaults(model_ema=True)
    parser.add_argument('--model-ema-decay', type=float, default=0.99996, help='')
    parser.add_argument('--model-ema-force-cpu', action='store_true', default=False, help='')

    # Optimizer parameters
    parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
                        help='Optimizer (default: "adamw"')
    parser.add_argument('--opt-eps', default=1e-8, type=float, metavar='EPSILON',
                        help='Optimizer Epsilon (default: 1e-8)')
    parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA',
                        help='Optimizer Betas (default: None, use opt default)')
    parser.add_argument('--clip-grad', type=float, default=None, metavar='NORM',
                        help='Clip gradient norm (default: None, no clipping)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                        help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', type=float, default=0.05,
                        help='weight decay (default: 0.05)')
    # Learning rate schedule parameters
    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER',
                        help='LR scheduler (default: "cosine"')
    parser.add_argument('--lr', type=float, default=5e-4, metavar='LR',
                        help='learning rate (default: 5e-4)')
    parser.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct',
                        help='learning rate noise on/off epoch percentages')
    parser.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT',
                        help='learning rate noise limit percent (default: 0.67)')
    parser.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV',
                        help='learning rate noise std-dev (default: 1.0)')
    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR',
                        help='warmup learning rate (default: 1e-6)')
    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')

    parser.add_argument('--decay-epochs', type=float, default=30, metavar='N',
                        help='epoch interval to decay LR')
    parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',
                        help='epochs to warmup LR, if scheduler supports')
    parser.add_argument('--cooldown-epochs', type=int, default=10, metavar='N',
                        help='epochs to cooldown LR at min_lr, after cyclic schedule ends')
    parser.add_argument('--patience-epochs', type=int, default=10, metavar='N',
                        help='patience epochs for Plateau LR scheduler (default: 10')
    parser.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE',
                        help='LR decay rate (default: 0.1)')

    # Augmentation parameters
    parser.add_argument('--color-jitter', type=float, default=0.4, metavar='PCT',
                        help='Color jitter factor (default: 0.4)')
    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
                        help='Use AutoAugment policy. "v0" or "original". " + \
                             "(default: rand-m9-mstd0.5-inc1)'),
    parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
    parser.add_argument('--train-interpolation', type=str, default='bicubic',
                        help='Training interpolation (random, bilinear, bicubic default: "bicubic")')

    parser.add_argument('--repeated-aug', action='store_true')
    parser.add_argument('--no-repeated-aug', action='store_false', dest='repeated_aug')
    parser.set_defaults(repeated_aug=True)

    # * Random Erase params
    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
                        help='Random erase prob (default: 0.25)')
    parser.add_argument('--remode', type=str, default='pixel',
                        help='Random erase mode (default: "pixel")')
    parser.add_argument('--recount', type=int, default=1,
                        help='Random erase count (default: 1)')
    parser.add_argument('--resplit', action='store_true', default=False,
                        help='Do not random erase first (clean) augmentation split')

    # * Mixup params
    parser.add_argument('--mixup', type=float, default=0.8,
                        help='mixup alpha, mixup enabled if > 0. (default: 0.8)')
    parser.add_argument('--cutmix', type=float, default=1.0,
                        help='cutmix alpha, cutmix enabled if > 0. (default: 1.0)')
    parser.add_argument('--cutmix-minmax', type=float, nargs='+', default=None,
                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
    parser.add_argument('--mixup-prob', type=float, default=1.0,
                        help='Probability of performing mixup or cutmix when either/both is enabled')
    parser.add_argument('--mixup-switch-prob', type=float, default=0.5,
                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
    parser.add_argument('--mixup-mode', type=str, default='batch',
                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')

    # Distillation parameters
    parser.add_argument('--teacher-model', default='regnety_160', type=str, metavar='MODEL',
                        help='Name of teacher model to train (default: "regnety_160"')
    parser.add_argument('--teacher-path', type=str, default='')
    parser.add_argument('--distillation-type', default='none', choices=['none', 'soft', 'hard'], type=str, help="")
    parser.add_argument('--distillation-alpha', default=0.5, type=float, help="")
    parser.add_argument('--distillation-tau', default=1.0, type=float, help="")

    # * Finetuning params
    parser.add_argument('--finetune', default='', help='finetune from checkpoint')

    # Dataset parameters
    parser.add_argument('--data-path', default='//project/ZHIHOUDATA/Code/RIDE-LongTailRecognition/data/ImageNet_LT/', type=str,
                        help='dataset path')
    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19'],
                        type=str, help='Image Net dataset path')
    parser.add_argument('--inat-category', default='name',
                        choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
                        type=str, help='semantic granularity')

    parser.add_argument('--output_dir', default='',
                        help='path where to save, empty for no saving')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--dist-eval', action='store_true', default=False, help='Enabling distributed evaluation')
    parser.add_argument('--num_workers', default=10, type=int)
    parser.add_argument('--pin-mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem',
                        help='')
    parser.set_defaults(pin_mem=True)

    # distributed training parameters
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--add_bt', default=0, type=int,
                        help='add_bt')
    parser.add_argument('--eval_global', default=0, type=int,
                        help='eval_global') # This is for evaluation with a mini-batch
    parser.add_argument('--dropout', default=0., type=float,
                        help='dropout')
    parser.add_argument('--exp_id', default="", type=str,
                        help='exp_id')

    # For BT
    parser.add_argument('--decay_drop_bt', default=0., type=float,
                        help='drop_patch')
    parser.add_argument('--drop_patch', default=0., type=float,
                        help='drop_patch')
    parser.add_argument('--shuffle_patch', action='store_true', )
    parser.add_argument('--no_fp16_bt', default=0, type=int,
                        help='no_fp16_bt')
    parser.add_argument('--num_heads', default=4, type=int,
                        help='num_heads')
    parser.add_argument('--start_idx', default=8, type=int,
                        help='start_idx: start bt layer index')
    parser.add_argument('--start_bt_epoch', default=0, type=int,
                        help='start_bt_epoch: start bt layer index')
    parser.add_argument('--insert_idx', action='append', type=int,
                        help='insert idx list')
    parser.add_argument('--all_patches', action='store_true',)
    parser.add_argument('--drop_path_bt', default=0., type=float)
    parser.add_argument('--not_cls_token', action='store_true'),
    parser.add_argument('--cls_token_only', action='store_true'),
    parser.add_argument('--shared_bt', default=1, type=int,),
    parser.add_argument('--empty_bt', default=0, type=int,),
    parser.add_argument('--add_norm_bt', default=0, type=int,), # This is to add BT in the last layer (normalization)
    parser.add_argument('--add_mlp_bt', default=0, type=int,),
    parser.add_argument('--mlp_enc', default=0, type=int,),
    parser.add_argument('--no_grad_bt', default=0, type=int,),
    parser.add_argument('--bt_decay', default=0., type=float,
                        help='drop_patch')
    parser.add_argument('--mlp_decay', default=0., type=float,
                       help='drop_patch')
    parser.add_argument('--bt_lr', default=0.5, type=float,)
    parser.add_argument('--balanced_softmax', action='store_true',) # This is for the experiments on Long-Tailed Recognition
    parser.add_argument('--bt_atten_drop', default=0.5, type=float)

    parser.add_argument('--skip_bt', action='store_true',)
    parser.add_argument('--pretrained', action='store_true',
                        help='pretrained')
    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
    return parser

def main(args):
    # torch.autograd.set_detect_anomaly(True)
    utils.init_distributed_mode(args)

    print(args)

    if args.distillation_type != 'none' and args.finetune and not args.eval:
        raise NotImplementedError("Finetuning with distillation not yet supported")

    device = torch.device(args.device)

    # fix the seed for reproducibility
    seed = args.seed + utils.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)
    # random.seed(seed)

    cudnn.benchmark = True

    dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
    dataset_val, _ = build_dataset(is_train=False, args=args)

    if True:  # args.distributed:
        num_tasks = utils.get_world_size()
        global_rank = utils.get_rank()
        if args.repeated_aug:
            sampler_train = RASampler(
                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
            )
        else:
            sampler_train = torch.utils.data.DistributedSampler(
                dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
            )
        if args.dist_eval:
            if len(dataset_val) % num_tasks != 0:
                print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
                      'This will slightly alter validation results as extra duplicate entries are added to achieve '
                      'equal num of samples per-process.')
            sampler_val = torch.utils.data.DistributedSampler(
                dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
        else:
            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
    else:
        sampler_train = torch.utils.data.RandomSampler(dataset_train)
        sampler_val = torch.utils.data.SequentialSampler(dataset_val)

    data_loader_train = torch.utils.data.DataLoader(
        dataset_train, sampler=sampler_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=True,
    )
    eval_batch_size = int(1.5 * args.batch_size)
    import os
    if os.environ['HOME'].endswith('root'):
        eval_batch_size = int(4 * args.batch_size)
    data_loader_val = torch.utils.data.DataLoader(
        dataset_val, sampler=sampler_val,
        batch_size=eval_batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=False
    )

    mixup_fn = None
    mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
    if mixup_active:
        mixup_fn = Mixup(
            mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
            prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
            label_smoothing=args.smoothing, num_classes=args.nb_classes)

    print(f"Creating model: {args.model}")
    model = create_model(
        args.model,
        pretrained=args.pretrained,
        num_classes=args.nb_classes,
        drop_rate=args.drop,
        drop_path_rate=args.drop_path,
        drop_block_rate=None,
    )
    linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 512.0
    args.lr = linear_scaled_lr
    args.bt_lr=linear_scaled_lr*args.bt_lr
    from optimizer import create_bt_optimizer
    model_o = model
    optimizer1 = create_bt_optimizer(args, model, bt_decay=args.bt_decay)
    lr_scheduler1, _ = create_scheduler(args, optimizer1)

    if args.add_bt:
        from bt import TransformerDecorator1, BlockBF, BlockWrap32
        encoder_global = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                               args=args, drop_path=args.drop_path_bt)

        if args.add_bt in [1]:
            # This is for the last layer
            res_blocks = []
            for i, block in enumerate(model.blocks):
                if args.no_fp16_bt and args.no_fp16_bt not in [4]:
                    # This is for ease nan
                    res_blocks.append(BlockWrap32(block, args.no_fp16_bt))
                else:
                    res_blocks.append(block)
            model.blocks = torch.nn.Sequential(*res_blocks)
            model.norm = torch.nn.Sequential(model.norm, encoder_global)
        elif args.add_bt in [3]:
            # for insert into a single layer
            res_blocks = []
            drop_rate_list = [args.drop_path_bt]*(len(model.blocks)+1)
            if args.insert_idx is not None and len(args.insert_idx) > 0:
                insert_list = args.insert_idx
            else:
                insert_list = list(range(args.start_idx, len(model.blocks)))
            for i, block in enumerate(model.blocks):
                if args.no_fp16_bt and args.no_fp16_bt not in [4]:
                    res_blocks.append(BlockWrap32(block, args.no_fp16_bt))
                else:
                    res_blocks.append(block)
                if i in insert_list:
                    if not args.shared_bt:
                        encoder_global = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                                           args=args, drop_path=drop_rate_list[i])
                    res_blocks.append(encoder_global)
            model.blocks = torch.nn.Sequential(*res_blocks)
            if args.add_norm_bt:
                if not args.shared_bt:
                    encoder_global = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                                        args=args, drop_path=args.drop_path_bt)
                model.norm = torch.nn.Sequential(model.norm, encoder_global)
        elif args.add_bt in [2]:
            # This is for multiple layers
            # 61, 62 is for half batch
            # model.norm = torch.nn.Sequential(model.norm, encoder_global)
            res_blocks = []
            first_enc = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                            args=args, first_layer=True, drop_path=args.drop_path_bt)
            old_enc = first_enc.encoder_layers
            del old_enc
            first_enc.encoder_layers = encoder_global.encoder_layers
            nums = 0
            if args.insert_idx is not None and len(args.insert_idx) > 0:
                insert_list = args.insert_idx
            else:
                insert_list = list(range(args.start_idx, len(model.blocks)))
            for i, block in enumerate(model.blocks):
                if args.no_fp16_bt and args.no_fp16_bt not in [4]:
                    res_blocks.append(BlockWrap32(block, args.no_fp16_bt))
                else:
                    res_blocks.append(block)
                if i in insert_list:
                    if insert_list[0] == i: # first layer
                        res_blocks.append(first_enc)
                        first_enc = None
                    else:
                        if not args.shared_bt:
                            encoder_global = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                                                    args=args, drop_path=args.drop_path_bt)
                        res_blocks.append(encoder_global)
            # model.norm = torch.nn.Identity()
            if args.add_norm_bt:
                if not args.shared_bt:
                    encoder_global = TransformerDecorator1(args.add_bt, model.num_features, args.eval_global,
                                                        args=args, drop_path=args.drop_path_bt)

                model.norm = torch.nn.Sequential(model.norm, encoder_global)
            model.blocks = torch.nn.Sequential(*res_blocks)

    if args.drop_patch != 0:
        # what's the relationship between drop_patch and shared prediction modules
        # This is not for the paper.
        res_blocks = []
        for i, block in enumerate(model.blocks):
            if i > args.start_idx:
                if args.drop_patch > 0:
                    from timm.models.layers import DropPath
                    block.drop_path = DropPath(args.drop_patch)
                    res_blocks.append(BlockBF(block, args.no_fp16_bt))
                else:
                    res_blocks.append(BlockBF(block, args.no_fp16_bt))
            else:
                res_blocks.append(block)
        model.blocks = torch.nn.Sequential(*res_blocks)
    print(model)
    if args.finetune:
        if args.finetune.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(
                args.finetune, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.finetune, map_location='cpu')

        checkpoint_model = checkpoint['model']
        state_dict = model.state_dict()
        for k in ['head.weight', 'head.bias', 'head_dist.weight', 'head_dist.bias']:
            if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
                print(f"Removing key {k} from pretrained checkpoint")
                del checkpoint_model[k]

        # interpolate position embedding
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
        # only the position tokens are interpolated
        pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
        pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
        pos_tokens = torch.nn.functional.interpolate(
            pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
        pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
        new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
        checkpoint_model['pos_embed'] = new_pos_embed

        model.load_state_dict(checkpoint_model, strict=False)
    if args.add_mlp_bt:
        from bt import MLPDecorder
        model.pre_logits = MLPDecorder(model.num_features, 4096, skip_mlp=args.skip_bt)

    print(model)
    model.to(device)

    model_ema = None
    if args.model_ema:
        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
        model_ema = ModelEma(
            model,
            decay=args.model_ema_decay,
            device='cpu' if args.model_ema_force_cpu else '',
            resume='')

    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True if args.skip_bt else False)
        model_without_ddp = model.module
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params:', n_parameters)

    # linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size() / 512.0
    args.lr = linear_scaled_lr
    from optimizer import create_bt_optimizer
    optimizer = create_bt_optimizer(args, model_without_ddp, bt_decay=args.bt_decay)
    loss_scaler = NativeScaler()

    lr_scheduler, _ = create_scheduler(args, optimizer)

    criterion = LabelSmoothingCrossEntropy()

    if mixup_active:
        # smoothing is handled with mixup label transform
        criterion = SoftTargetCrossEntropy()
    elif args.smoothing:
        criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
    else:
        criterion = torch.nn.CrossEntropyLoss()

    teacher_model = None
    if args.distillation_type != 'none':
        assert args.teacher_path, 'need to specify teacher-path when using distillation'
        print(f"Creating teacher model: {args.teacher_model}")
        teacher_model = create_model(
            args.teacher_model,
            pretrained=args.pretrained,
            num_classes=args.nb_classes,
            global_pool='avg',
        )
        if args.teacher_path.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(
                args.teacher_path, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.teacher_path, map_location='cpu')
        teacher_model.load_state_dict(checkpoint['model'])
        teacher_model.to(device)
        teacher_model.eval()

    # wrap the criterion in our custom DistillationLoss, which
    # just dispatches to the original criterion if args.distillation_type is 'none'
    criterion = DistillationLoss(
        criterion, teacher_model, args.distillation_type, args.distillation_alpha, args.distillation_tau
    )
    if args.balanced_softmax:
        from losses import BalancedSoftmax
        criterion = BalancedSoftmax('./cls_freq/ImageNet_LT.json')
    output_dir = Path(args.output_dir)
    if args.resume:
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(
                args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        if args.skip_bt:
            model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
        else:
            model_without_ddp.load_state_dict(checkpoint['model'], strict=True)
        if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
            optimizer.load_state_dict(checkpoint['optimizer'])
            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
            args.start_epoch = checkpoint['epoch'] + 1
            if args.model_ema:
                utils._load_checkpoint_for_ema(model_ema, checkpoint['model_ema'])
            if 'scaler' in checkpoint:
                loss_scaler.load_state_dict(checkpoint['scaler'])

    if args.skip_bt:
        optimizer.param_groups.pop(1)
        if args.add_mlp_bt:
            logits_list = [item for item in model_without_ddp.pre_logits.parameters()]
            print(len(optimizer.param_groups), len(logits_list))
            new_list = []
            if len(optimizer.param_groups) == 1:
                for item in optimizer.param_groups[0]['params']:
                    is_logits=False
                    for item1 in logits_list:
                        if item.size == item1.size and torch.all(torch.eq(item, item1)):
                            is_logits = True
                            break
                    if not is_logits:
                        new_list.append(item)
                optimizer.param_groups[0]['params'] = new_list
            else:
                for item in optimizer.param_groups[1]['params']:
                    is_logits=False
                    for item1 in logits_list:
                        if item.size == item1.size and torch.all(torch.eq(item, item1)):
                            is_logits = True
                            break
                    if not is_logits:
                        new_list.append(item)
                optimizer.param_groups[1]['params'] = new_list

    if args.eval:
        # test_stats = evaluate(data_loader_val, model, device)
        training_labels = dataset_train.targets
        training_labels = np.array(training_labels).astype(int)
        from engine import validate
        validate(data_loader_val, training_labels, model, device, args.nb_classes, utils.get_rank())

        print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
        return

    print(f"Start training for {args.epochs} epochs")
    start_time = time.time()
    max_accuracy = 0.0
    opt2=optimizer
    model2 = model
    lr_sched = lr_scheduler
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            data_loader_train.sampler.set_epoch(epoch)

        if epoch < args.start_bt_epoch:
            for item in model.modules():
                if isinstance(item, TransformerDecorator1):
                    item.eval()
                    for n,p in item.named_parameters():
                        p.requires_grad=False
            optimizer=optimizer1
            model = model_o
            lr_scheduler = lr_scheduler1
        elif epoch == args.start_bt_epoch:

            for n,p in model.named_parameters():
                p.requires_grad=True
            optimizer=opt2
            model = model2
            lr_scheduler = lr_sched
        model.train(args.finetune == '')

        train_stats = train_one_epoch(
            model, criterion, data_loader_train,
            optimizer, device, epoch, loss_scaler,
            args.clip_grad, model_ema, mixup_fn,
            set_training_mode=args.finetune == '',  # keep in eval mode during finetuning
            no_fp16_bt=args.no_fp16_bt, model_without_ddp = model_without_ddp,
            num_classes=args.nb_classes
        )
        if epoch < args.start_bt_epoch:
            lr_sched.step(epoch)
        lr_scheduler.step(epoch)
        if args.output_dir:
            checkpoint_paths = [output_dir / str(args.exp_id+'checkpoint.pth')]
            for checkpoint_path in checkpoint_paths:
                utils.save_on_master({
                    'model': model_without_ddp.state_dict(),
                    'optimizer': optimizer.state_dict() if not args.skip_bt else 0,
                    'lr_scheduler': lr_scheduler.state_dict(),
                    'epoch': epoch,
                    'model_ema': get_state_dict(model_ema),
                    'scaler': loss_scaler.state_dict(),
                    'args': args,
                }, checkpoint_path)

        test_stats = evaluate(data_loader_val, model, device)
        print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
        max_accuracy = max(max_accuracy, test_stats["acc1"])
        print(f'Max accuracy: {max_accuracy:.2f}%')

        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                     **{f'test_{k}': v for k, v in test_stats.items()},
                     'epoch': epoch,
                     'n_parameters': n_parameters}

        if args.output_dir and utils.is_main_process():
            with (output_dir / str(args.exp_id+"log.txt")).open("a") as f:
                f.write(json.dumps(log_stats) + "\n")

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))


if __name__ == '__main__':
    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)
